import sys
import json
import base64
import requests
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QTabWidget, 
                              QVBoxLayout, QHBoxLayout, QGroupBox, QLabel, 
                              QPushButton, QTextEdit, QFileDialog, QDialog,
                              QLineEdit, QComboBox, QDialogButtonBox, QFormLayout,)
from PySide6.QtCore import Qt, QByteArray, QBuffer, QIODevice, QThread, Signal
from PySide6.QtGui import QPixmap, QImage, QFont, QIcon

class ConfigDialog(QDialog):
    def __init__(self, config, parent=None):
        super().__init__(parent)
        self.setWindowTitle("配置设置")
        self.setWindowIcon(QIcon("setting.svg"))
        self.setGeometry(200, 200, 500, 400)
        
        self.config = config
        
        self.create_ui()
    
    def create_ui(self):
        layout = QVBoxLayout()
        self.setLayout(layout)
        
        tabs = QTabWidget()
        layout.addWidget(tabs)
        
        # 图像转Markdown配置
        ocr_tab = QWidget()
        ocr_layout = QFormLayout(ocr_tab)
        
        self.ocr_api_key = QLineEdit(self.config['image_to_markdown']['api_key'])
        self.ocr_api_key.setPlaceholderText("输入OpenAI API密钥")
        
        self.ocr_model = QLineEdit(self.config['image_to_markdown']['model'])
        self.ocr_model.setPlaceholderText("例如: gpt-4-vision-preview")
        
        self.ocr_base_url = QLineEdit(self.config['image_to_markdown']['base_url'])
        self.ocr_base_url.setPlaceholderText("例如: https://api.openai.com/v1")
        
        self.ocr_prompt = QTextEdit()
        self.ocr_prompt.setText(self.config['image_to_markdown']['prompt'])
        self.ocr_prompt.setPlaceholderText("输入图像转换的提示词")
        self.ocr_prompt.setMaximumHeight(100)
        
        ocr_layout.addRow("API密钥:", self.ocr_api_key)
        ocr_layout.addRow("模型:", self.ocr_model)
        ocr_layout.addRow("基础URL:", self.ocr_base_url)
        ocr_layout.addRow("提示词:", self.ocr_prompt)
        
        tabs.addTab(ocr_tab, "图像转Markdown")
        
        # 翻译配置
        trans_tab = QWidget()
        trans_layout = QFormLayout(trans_tab)
        
        self.trans_api_key = QLineEdit(self.config['translation']['api_key'])
        self.trans_api_key.setPlaceholderText("输入OpenAI API密钥")
        
        self.trans_model = QLineEdit(self.config['translation']['model'])
        self.trans_model.setPlaceholderText("例如: gpt-4")
        
        self.trans_base_url = QLineEdit(self.config['translation']['base_url'])
        self.trans_base_url.setPlaceholderText("例如: https://api.openai.com/v1")
        
        self.trans_auto_prompt = QTextEdit()
        self.trans_auto_prompt.setText(self.config['translation']['prompt_auto'])
        self.trans_auto_prompt.setPlaceholderText("输入自动检测语言时的提示词")
        self.trans_auto_prompt.setMaximumHeight(80)
        
        trans_layout.addRow("API密钥:", self.trans_api_key)
        trans_layout.addRow("模型:", self.trans_model)
        trans_layout.addRow("基础URL:", self.trans_base_url)
        trans_layout.addRow("自动检测提示词:", self.trans_auto_prompt)
        
        tabs.addTab(trans_tab, "文本翻译")
        
        # 按钮
        button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addWidget(button_box)
    
    def get_config(self):
        # 返回更新后的配置
        return {
            "image_to_markdown": {
                "api_key": self.ocr_api_key.text(),
                "model": self.ocr_model.text(),
                "base_url": self.ocr_base_url.text(),
                "prompt": self.ocr_prompt.toPlainText()
            },
            "translation": {
                "api_key": self.trans_api_key.text(),
                "model": self.trans_model.text(),
                "base_url": self.trans_base_url.text(),
                "prompt_auto": self.trans_auto_prompt.toPlainText(),
                "prompt_map": self.config['translation']['prompt_map']  # 内部映射保持不变
            }
        }

class MainWindow(QMainWindow):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.setWindowTitle("AI工具箱")
        self.setWindowIcon(QIcon("app_icon.png"))  # 添加这行设置窗口图标
        self.setGeometry(100, 100, 900, 700)
        self.setMinimumSize(800, 600)
        
        # 先创建标签页再创建设置按钮
        self.tabs = QTabWidget()  # 提前初始化
        self.tabs.setFont(QFont("Microsoft YaHei UI", 10))
        self.tabs.addTab(self.create_image_tab(), "图像转Markdown")
        self.tabs.addTab(self.create_translation_tab(), "文本翻译")
        
        self.create_settings_button()  # 在标签页之后初始化
        
        self.setCentralWidget(self.tabs)
        self.apply_styles()

        # 翻译语言选择
        self.translation_lang = "auto"  # 默认自动检测
        
    def create_settings_button(self):
        """创建标签页右上角的设置按钮"""
        settings_button = QPushButton("配置")
        settings_button.setToolTip("设置")
        settings_button.setStyleSheet("""
            QPushButton {
                background-color: #404040;  /* 更改为深灰色 */
                color: #A0A0A0;           /* 文字颜色改为浅灰色 */
                border: 1px solid #505050;
                border-radius: 4px;
                padding: 5px;
                min-width: 60px;
                margin: 2px;             /* 添加边距避免贴边 */
            }
            QPushButton:hover {
                background-color: #505050;
                border-color: #606060;
            }
        """)
        settings_button.clicked.connect(self.open_config_dialog)
        
        # 将按钮添加到标签页右上角
        self.tabs.setCornerWidget(settings_button, Qt.TopRightCorner)
    
    def open_config_dialog(self):
        dialog = ConfigDialog(self.config, self)
        if dialog.exec() == QDialog.Accepted:
            self.config = dialog.get_config()
            self.save_config()
    
    def save_config(self):
        try:
            with open("config.json", "w", encoding="utf-8") as f:
                json.dump(self.config, f, indent=4, ensure_ascii=False)
            print("配置已保存到 config.json")
        except Exception as e:
            print(f"保存配置失败: {e}")
    
    def apply_styles(self):
        self.setStyleSheet("""
            QMainWindow {
                background-color: #2D2D30;
            }
            QTabWidget {
                background-color: #252526;
                color: #D4D4D4;
                border: none;
                font-size: 12pt;
            }
            QTabBar::tab {
                background-color: #1E1E1E;
                color: #CCCCCC;
                padding: 10px 20px;
                border-top-left-radius: 5px;
                border-top-right-radius: 5px;
                margin-right: 2px;
            }
            QTabBar::tab:selected {
                background-color: #0078D7;
                color: white;
                font-weight: bold;
            }
            QGroupBox {
                border: 1px solid #3F3F46;
                border-radius: 5px;
                margin-top: 1em;
                padding-top: 15px;
                font-weight: bold;
                color: #9CDCFE;
            }
            QPushButton {
                background-color: #0078D7;
                color: white;
                border: none;
                border-radius: 4px;
                padding: 8px 16px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #3A96DD;
            }
            QTextEdit {
                background-color: #1E1E1E;
                color: #D4D4D4;
                border: 1px solid #3F3F46;
                border-radius: 4px;
                padding: 8px;
                font-size: 11pt;
            }
            QLabel {
                color: #CCCCCC;
            }
            .dropArea {
                border: 2px dashed #0078D7;
                border-radius: 10px;
                padding: 20px;
                background-color: rgba(30, 90, 150, 0.1);
                text-align: center;
            }
            .title {
                font-size: 18pt;
                font-weight: bold;
                color: #4EC9B0;
                margin-bottom: 15px;
            }
            QComboBox {
                background-color: #1E1E1E;
                color: #D4D4D4;
                border: 1px solid #3F3F46;
                border-radius: 4px;
                padding: 5px;
            }
        """)
    
    def create_image_tab(self):
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 标题
        title = QLabel("图像转Markdown工具")
        title.setProperty("class", "title")
        title.setAlignment(Qt.AlignCenter)
        layout.addWidget(title)
        
        # 输入区域
        input_group = QGroupBox("图片输入")
        input_layout = QVBoxLayout()
        
        self.drop_label = QLabel("拖拽图片到此处或点击上传按钮")
        self.drop_label.setAlignment(Qt.AlignCenter)
        self.drop_label.setProperty("class", "dropArea")
        self.drop_label.setMinimumHeight(200)
        self.drop_label.setAcceptDrops(True)
        self.drop_label.setCursor(Qt.PointingHandCursor)
        
        btn_layout = QHBoxLayout()
        btn_upload = QPushButton("上传图片")
        btn_paste = QPushButton("粘贴图片")
        btn_clear = QPushButton("清除")
        
        btn_upload.clicked.connect(self.upload_image)
        btn_paste.clicked.connect(self.paste_image)
        btn_clear.clicked.connect(self.clear_image)
        
        btn_layout.addWidget(btn_upload)
        btn_layout.addWidget(btn_paste)
        btn_layout.addWidget(btn_clear)
        
        input_layout.addWidget(self.drop_label)
        input_layout.addLayout(btn_layout)
        input_group.setLayout(input_layout)
        layout.addWidget(input_group)
        
        # 输出区域
        output_group = QGroupBox("Markdown输出")
        output_layout = QVBoxLayout()
        
        self.output_text = QTextEdit()
        self.output_text.setPlaceholderText("转换结果将显示在这里...")
        self.output_text.setReadOnly(False)
        
        btn_layout_output = QHBoxLayout()
        btn_convert = QPushButton("开始转换")
        btn_convert.setObjectName("btn_convert")
        
        btn_copy = QPushButton("复制到剪贴板")
        btn_convert.clicked.connect(self.convert_image)
        btn_copy.clicked.connect(self.copy_markdown)
        
        btn_layout_output.addWidget(btn_convert)
        btn_layout_output.addWidget(btn_copy)
        
        output_layout.addWidget(self.output_text)
        output_layout.addLayout(btn_layout_output)
        output_group.setLayout(output_layout)
        layout.addWidget(output_group)
        
        return tab
    
    def create_translation_tab(self):
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 标题
        title = QLabel("文本翻译工具")
        title.setProperty("class", "title")
        title.setAlignment(Qt.AlignCenter)
        layout.addWidget(title)
        
        # 语言选择
        lang_layout = QHBoxLayout()
        lang_label = QLabel("翻译方向:")
        self.lang_combo = QComboBox()
        self.lang_combo.addItem("自动检测", "auto")
        self.lang_combo.addItem("英文 → 中文", "en2cn")
        self.lang_combo.addItem("中文 → 英文", "cn2en")
        self.lang_combo.addItem("英文 → 日文", "en2jp")
        self.lang_combo.addItem("日文 → 中文", "jp2cn")
        self.lang_combo.setCurrentIndex(0)  # 默认自动检测
        
        lang_layout.addWidget(lang_label)
        lang_layout.addWidget(self.lang_combo)
        lang_layout.addStretch()
        layout.addLayout(lang_layout)
        
        # 输入区域
        input_group = QGroupBox("输入文本")
        input_layout = QVBoxLayout()
        
        self.translate_input = QTextEdit()
        self.translate_input.setPlaceholderText("请输入要翻译的内容...")
        
        input_layout.addWidget(self.translate_input)
        input_group.setLayout(input_layout)
        layout.addWidget(input_group)
        
        # 输出区域
        output_group = QGroupBox("翻译结果")
        output_layout = QVBoxLayout()
        
        self.translate_output = QTextEdit()
        self.translate_output.setPlaceholderText("翻译结果将显示在这里...")
        
        btn_layout = QHBoxLayout()
        btn_translate = QPushButton("开始翻译")
        btn_translate.setObjectName("btn_translate")
        btn_copy_trans = QPushButton("复制到剪贴板")
        
        btn_translate.clicked.connect(self.translate_text)
        btn_copy_trans.clicked.connect(self.copy_translation)
        
        btn_layout.addWidget(btn_translate)
        btn_layout.addWidget(btn_copy_trans)
        
        output_layout.addWidget(self.translate_output)
        output_layout.addLayout(btn_layout)
        
        output_group.setLayout(output_layout)
        layout.addWidget(output_group)
        
        return tab
    
    def get_translation_prompt(self):
        """根据选择的语言方向获取提示词"""
        lang_option = self.lang_combo.currentData()
        
        if lang_option == "auto":
            # 使用用户自定义的auto提示词
            return self.config['translation']['prompt_auto']
        else:
            # 使用预定义的提示词
            prompts = self.config['translation']['prompt_map']
            return prompts.get(lang_option, self.config['translation']['prompt_auto'])
    
    def dragEnterEvent(self, event):
        if event.mimeData().hasImage() or event.mimeData().hasUrls():
            event.acceptProposedAction()
    
    def dropEvent(self, event):
        mime_data = event.mimeData()
        self.handle_drop(mime_data)
    
    def handle_drop(self, mime_data):
        """处理拖拽的数据"""
        # 检查是否有图像数据
        if mime_data.hasImage():
            image = mime_data.imageData()
            if not image.isNull():
                self.show_image(image)
                return
        
        # 检查是否有URL（文件路径）
        if mime_data.hasUrls():
            url = mime_data.urls()[0]
            if url.isLocalFile():
                file_path = url.toLocalFile()
                # 检查文件扩展名
                if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')):
                    image = QImage(file_path)
                    if not image.isNull():
                        self.show_image(image)
                        return
        
        self.drop_label.setText("无效的图像文件")
    
    def show_image(self, image):
        """显示选中的图片"""
        self.current_image = image
        pixmap = QPixmap.fromImage(image.scaled(400, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation))
        self.drop_label.setPixmap(pixmap)
        self.drop_label.setText("")
    
    def upload_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, 
            "选择图片", 
            "", 
            "图片文件 (*.png *.jpg *.jpeg *.bmp *.gif *.webp)"
        )
        if file_path:
            image = QImage(file_path)
            if not image.isNull():
                self.show_image(image)
            else:
                self.drop_label.setText("无法加载图像")
    
    def paste_image(self):
        """从剪贴板粘贴图像 (更可靠的版本)"""
        clipboard = QApplication.clipboard()
        mime_data = clipboard.mimeData()
        
        # 检查是否可以直接获取图像
        if mime_data.hasImage():
            image = clipboard.image()
            if not image.isNull():
                self.show_image(image)
                return
                
        # 检查是否可以通过URL获取图像
        if mime_data.hasUrls():
            for url in mime_data.urls():
                if url.isLocalFile():
                    file_path = url.toLocalFile()
                    if any(file_path.lower().endswith(ext) for ext in ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')):
                        image = QImage(file_path)
                        if not image.isNull():
                            self.show_image(image)
                            return
                            
        # 检查是否有文件路径作为文本
        if mime_data.hasText():
            text = clipboard.text().strip()
            if any(text.lower().endswith(ext) for ext in ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')):
                image = QImage(text)
                if not image.isNull():
                    self.show_image(image)
                    return
        
        # 尝试直接创建QImage
        if mime_data.formats():
            for format in mime_data.formats():
                if "image" in format.lower():
                    data = mime_data.data(format)
                    image = QImage()
                    if image.loadFromData(data, format):
                        if not image.isNull():
                            self.show_image(image)
                            return
        
        # 如果所有方法都失败
        self.drop_label.setText("剪贴板中没有图片或无法识别")
    
    def clear_image(self):
        """清除当前图像和结果"""
        self.drop_label.clear()
        self.drop_label.setText("拖拽图片到此处或点击上传按钮")
        self.output_text.clear()
        if hasattr(self, 'current_image'):
            del self.current_image
    
    class NetworkWorker(QThread):
        finished = Signal(object)
        error = Signal(str)
    
        def __init__(self, url, headers, payload):
            super().__init__()
            self.url = url
            self.headers = headers
            self.payload = payload
    
        def run(self):
            try:
                response = requests.post(
                    self.url,
                    headers=self.headers,
                    json=self.payload,
                    timeout=60
                )
                response.raise_for_status()
                self.finished.emit(response.json())
            except Exception as e:
                self.error.emit(str(e))

    def convert_image(self):
        """将图像转换为Markdown文本"""
        if not hasattr(self, 'current_image'):
            self.output_text.setText("请先输入图片")
            return
    
        # 将图像转换为Base64编码
        ba = QByteArray()
        buffer = QBuffer(ba)
        buffer.open(QIODevice.WriteOnly)
        # 保存为PNG格式
        success = self.current_image.save(buffer, "PNG")
        if not success:
            self.output_text.setText("无法处理图像")
            return
        
        base64_image = base64.b64encode(ba.data()).decode('utf-8')
        
        config = self.config['image_to_markdown']
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config['api_key']}"
        }
        
        payload = {
            "model": config['model'],
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": config['prompt']},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
                    ]
                }
            ],
            "max_tokens": 2000
        }
        
        self.output_text.setText("正在转换中，请稍候...")
        QApplication.processEvents()  # 更新UI
        
        # 替换原来的网络请求部分
        self.output_text.setText("正在转换中，请稍候...")
        self.findChild(QPushButton, "btn_convert").setEnabled(False)  # 禁用转换按钮
        
        self.worker = self.NetworkWorker(
            f"{config['base_url']}/chat/completions",
            headers,
            payload
        )
        self.worker.finished.connect(self.handle_conversion_success)
        self.worker.error.connect(self.handle_conversion_error)
        self.worker.start()
    
    def handle_conversion_success(self, result):
        markdown_text = result['choices'][0]['message']['content']
        # 清理多余的反引号和标签
        markdown_text = markdown_text.strip().strip('`').strip()
        if markdown_text.lower().startswith('markdown'):
            markdown_text = markdown_text[8:].lstrip()
        self.output_text.setPlainText(markdown_text)
        self.findChild(QPushButton, "btn_convert").setEnabled(True)  # 恢复按钮
    
    def handle_conversion_error(self, error_msg):
        self.output_text.setText(f"转换失败: {error_msg}")
        self.findChild(QPushButton, "btn_convert").setEnabled(True)
    
    class TranslationWorker(QThread):
        finished = Signal(object)
        error = Signal(str)
        
        def __init__(self, url, headers, payload):
            super().__init__()
            self.url = url
            self.headers = headers
            self.payload = payload

        def run(self):
            try:
                response = requests.post(
                    self.url,
                    headers=self.headers,
                    json=self.payload,
                    timeout=30
                )
                response.raise_for_status()
                self.finished.emit(response.json())
            except Exception as e:
                self.error.emit(str(e))

    def translate_text(self):
        """翻译文本"""
        text = self.translate_input.toPlainText().strip()
        if not text:
            self.translate_output.setText("请输入要翻译的内容")
            return
        
        # 禁用翻译按钮
        btn_translate = self.findChild(QPushButton, "btn_translate")
        if btn_translate:
            btn_translate.setEnabled(False)

        # 获取配置和设置请求参数
        prompt = self.get_translation_prompt()
        config = self.config['translation']
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config['api_key']}"
        }
        
        payload = {
            "model": config['model'],
            "messages": [
                {"role": "system", "content": prompt},
                {"role": "user", "content": text}
            ],
            "max_tokens": 2000
        }
        
        # 显示等待提示
        self.translate_output.setText("正在翻译中，请稍候...")
        QApplication.processEvents()

        # 创建并启动工作线程
        self.trans_worker = self.TranslationWorker(
            f"{config['base_url']}/chat/completions",
            headers,
            payload
        )
        self.trans_worker.finished.connect(self.handle_translation_success)
        self.trans_worker.error.connect(self.handle_translation_error)
        self.trans_worker.start()

    def handle_translation_success(self, result):
        """处理翻译成功结果"""
        try:
            translation = result['choices'][0]['message']['content'].strip()
            self.translate_output.setPlainText(translation)
        except KeyError as e:
            self.translate_output.setText(f"解析响应出错: {str(e)}")
        finally:
            self.findChild(QPushButton, "btn_translate").setEnabled(True)

    def handle_translation_error(self, error_msg):
        """处理翻译错误"""
        self.translate_output.setText(f"翻译失败: {error_msg}")
        self.findChild(QPushButton, "btn_translate").setEnabled(True)

    def copy_markdown(self):
        QApplication.clipboard().setText(self.output_text.toPlainText())
    
    def copy_translation(self):
        QApplication.clipboard().setText(self.translate_output.toPlainText())

def load_config():
    """加载配置文件"""
    try:
        with open("config.json", "r", encoding="utf-8") as f:
            config = json.load(f)
        
        # 添加默认值以兼容旧版配置
        if 'prompt_auto' not in config['translation']:
            config['translation']['prompt_auto'] = "你是一个专业翻译助手，能够识别原文语言并准确翻译成目标语言："
        
        if 'prompt_map' not in config['translation']:
            config['translation']['prompt_map'] = {
                "en2cn": "你是一个专业翻译助手，将以下英文内容准确、流畅地翻译成简体中文：",
                "cn2en": "你是一个专业翻译助手，将以下中文内容准确、流畅地翻译成自然英文：",
                "en2jp": "You are a professional translator, please accurately translate the following English content into natural Japanese:",
                "jp2cn": "你是一个专业翻译助手，将以下日文内容准确、流畅地翻译成简体中文："
            }
        
        return config
    except FileNotFoundError:
        print("配置文件 config.json 未找到，使用默认配置")
    except json.JSONDecodeError:
        print("配置文件格式错误，使用默认配置")
    except Exception as e:
        print(f"配置文件加载失败: {e}")
    
    # 默认配置
    return {
        "image_to_markdown": {
            "api_key": "sk-xxxxxxxxxxxxxxxxxxxxxx",
            "model": "gpt-4-vision-preview",
            "base_url": "https://api.openai.com/v1",
            "prompt": "将图片中的内容准确转换为Markdown格式，包含数学公式的请使用LaTeX格式输出"
        },
        "translation": {
            "api_key": "sk-xxxxxxxxxxxxxxxxxxxxxx",
            "model": "gpt-4",
            "base_url": "https://api.openai.com/v1",
            "prompt_auto": "你是一个专业翻译助手，能够识别原文语言并准确翻译成目标语言：",
            "prompt_map": {
                "en2cn": "你是一个专业翻译助手，将以下英文内容准确、流畅地翻译成简体中文：",
                "cn2en": "你是一个专业翻译助手，将以下中文内容准确、流畅地翻译成自然英文：",
                "en2jp": "You are a professional translator, please accurately translate the following English content into natural Japanese:",
                "jp2cn": "你是一个专业翻译助手，将以下日文内容准确、流畅地翻译成简体中文："
            }
        }
    }

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setFont(QFont("Microsoft YaHei UI", 10))
    app.setStyle("Fusion")
    
    config = load_config()
    window = MainWindow(config)
    window.show()
    
    sys.exit(app.exec())